#!/usr/bin/env python


import types
import curses


  ############################################################
  ##       Start of cursesPlotter class definitions         ##
  ############################################################


STYLE = "."


  ##############
  ############
  ##
  ##  Figure:
  ##    Controls sizing
  ##    and placement
  ##
  ##########################
  ########################

class Figure:

    def __init__(self, win, ypanels=1, xpanels=1):
        self.win = win
        self.maxy, self.maxx = self.win.getmaxyx()
        self.ypanels = ypanels
        self.xpanels = xpanels
        self.panels = []
        for i in range(self.ypanels):
            self.panels.append([])
            for j in range(self.xpanels):
                self.panels[i].append(None)

    def plot(self, i, j, xdata=[0], ydata=[0], name=None):
        if not name: name="plot(%i,%i)" %(i,j)
        try:
            self.panels[i][j] = Plot(xdata, ydata, name)
            self.panels[i][j].panel = (i,j)
            return self.panels[i][j]
        except IndexError:
            raise PanelError(i,j)

    def refresh(self):
        self.maxy, self.maxx = self.win.getmaxyx()
        height =  int(self.maxy/float(self.ypanels))
        width =  int(self.maxx/float(self.xpanels))
        for i in range(self.ypanels):
            for j in range(self.xpanels):
                row = i * height
                col = j * width
                self.panels[i][j].win = self.win.derwin(height,width,row,col)
                self.panels[i][j].get_limits()
                self.panels[i][j].update()

    def get_plot(self, name=None, pos=(None,None)):
        for i in range(self.ypanels):
            for j in range(self.xpanels):
                if self.panels[i][j].name == name:
                    return self.panels[i][j]
                elif self.panels[i][j].win.enclose(*pos):
                    return self.panels[i][j]
        return None


  ##############
  ############
  ##
  ##  Plot:
  ##    Does actual plot-
  ##    ting of data
  ##
  ##########################
  ########################

class Plot:

    def __init__(self, xdata, ydata, name):
        self.allowable_types = set(( types.NoneType,
                                     types.IntType,
                                     types.TupleType,
                                     types.ListType ))
        if ( type(xdata) or type(ydata) ) not in self.allowable_types:
            raise TypeError, "Data type not currently support in Plot."
        else:
            self.xdata = xdata
            self.ydata = ydata
            self.name = name
            self.limits = {}
            self.preset_limits = False

    def update(self):
        self.win.erase()
        for p in self.points:
            self.win.addstr(self.maxy-p[0]-1,p[1],STYLE)
        self.win.border()
        self.win.addstr(0,int((self.maxx-len(self.name))/2.0),self.name)
        self.win.addstr(1,0,"%.2f" %self.limits["top"])
        if self.limits["btm"]==self.limits["lft"]==0:
            self.win.addstr(self.maxy-1,0,"0")
        else:
            self.win.addstr(self.maxy-2,0,"%.2f" %self.limits["btm"])
            self.win.addstr(self.maxy-1,2,"%.2f" %self.limits["lft"])
        self.win.addstr(self.maxy-1,self.maxx-6,"%.2f" %self.limits["rgt"])

    def set_limits(self, limits):
        for lim in limits:
            self.limits[lim] = limits[lim]
        self.preset_limits = True

    def get_limits(self):
        self.maxy, self.maxx = self.win.getmaxyx()
        if not self.preset_limits:
            self.limits = { "top": max(self.ydata),
                            "btm": min(self.ydata),
                            "rgt": max(self.xdata),
                            "lft": min(self.xdata) }
        self.xrange = self.limits["rgt"] - self.limits["lft"]
        self.yrange = self.limits["top"] - self.limits["btm"]
        self.points = []
        for x in range(self.maxx-2):
            for y in range(self.maxy-2):
                btm = self.limits["btm"] + y*(self.yrange/(self.maxy-2))
                top = self.limits["btm"] + (y+1)*(self.yrange/(self.maxy-2))
                lft = self.limits["lft"] + x*(self.xrange/(self.maxx-2))
                rgt = self.limits["lft"] + (x+1)*(self.xrange/(self.maxx-2))
                for i in range(len(self.ydata)):
                    if btm <= self.ydata[i] < top:
                        if lft <= self.xdata[i] < rgt:
                            self.points.append([y,x])
                    


  ##############
  ############
  ##
  ##  Exceptions:
  ##    handle plot-
  ##    ter erros
  ##
  ##########################
  ########################

class PanelError(Exception):
    def __init__(self, i, j):
        self.i = i
        self.j = j
    def __str__(self):
        return "Cannot plot to panel %i,%i.\
        It does not exist" %(self.i,self.j)



  ############################################################
  ##        End of cursesPlotter class definitions          ##
  ############################################################


from math import pi
from time import sleep
from random import uniform
from threading import Thread


class DataTracker(Thread):

    def __init__(self, example):
        self.example = example
        self.mainwin = self.example.get_win()
        self.quit = False
        self.pause = True
        self.para = (1.,1.)
        Thread.__init__(self)

    def run(self):
        while not self.quit:
            self.mainwin = self.example.get_win()
            if not self.pause:
                self.mainwin.erase()
                self.example.refresh_data(*self.para)
                self.example.update_fig()
                self.mainwin.refresh()
                sleep(1)


  ##############
  ############
  ##
  ##  ExampePlotter:
  ##    example class
  ##    for plotting
  ##
  ##########################
  ########################

class ExamplePlotter:

    def __init__(self, stdscr, ypanels, xpanels):
        self.quit = False
        self.stdscr = stdscr
        self.ypanels = ypanels
        self.xpanels = xpanels
        self.status = ("Welcome to the cursesPlotter example! Press 'p' to start the data tracker, or 'h' for a help screen",0)
        self.resize()

        xdata = [float(i) for i in range(-16,16)]
        ydata = [uniform(-pi/4.,pi/4.) for i in range(-16,16)]

        self.fig = Figure(self.mainwin,self.ypanels,self.xpanels)
        for a in range(self.ypanels):
            for b in range(self.xpanels):
                p = self.fig.plot(a,b,xdata,ydata)
                p.set_limits({"top":pi, "btm":-pi, "lft":-16.0, "rgt":16.0})
        self.fig.refresh()

        self.update_fig()


    def refresh_data(self,a,b):
        xdata = [float(i) for i in range(-16,16)]
        ydata = [uniform(-pi/a,pi/b) for i in range(-16,16)]
        for a in range(self.ypanels):
            for b in range(self.xpanels):
                self.fig.panels[a][b].ydata = ydata
                self.fig.panels[a][b].xdata = xdata


    def get_win(self):
        return self.mainwin


    def resize(self):
        self.maxy, self.maxx = self.stdscr.getmaxyx()
        self.centrx = int(round(self.maxx/2.0))
        self.centry = int(round(self.maxy/2.0))
        
        self.mainwin = curses.newwin(self.maxy-3, self.maxx, 0,0)
        self.statuswin = curses.newwin(3,self.maxx,self.maxy-3,0)


    def update_status(self, status=None):
        if not status:
            status = self.status
        self.status = status
        self.statuswin.erase()
        if status[1]==-1:
            self.statuswin.bkgd(" ",curses.A_STANDOUT)
        else:
            self.statuswin.bkgd(" ",0)
        self.statuswin.border()
        self.statuswin.addstr(1,self.centrx-int(len(status[0])/2.0),status[0])
        self.statuswin.refresh()


    def update_fig(self):
        self.fig.win = self.mainwin
        self.fig.refresh()
        self.mainwin.refresh()
        

    def show_help(self):
        self.help_str = (
            "<><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><>",
            "cursesPlotter accepts the following commands:",
            "q        -- quit the main loop",
            "h        -- show this help screen",
            "<space>  -- change plotting style",
            "p        -- redraw the current plot",
            " ",
            "screen is %i by %i" %(self.maxx,self.maxy),
            "<><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><><>"
            )
        for i in range(len(self.help_str)):
            self.mainwin.addstr(self.centry-int(len(self.help_str)/2.0)+i,
                                self.centrx-int(len(self.help_str[0])/2.0),
                                self.help_str[i])


def main(stdscr):
    #curses.mousemask(curses.ALL_MOUSE_EVENTS)
    curses.curs_set(0)

    quit = False

    example = ExamplePlotter(stdscr, 4, 2)
    example.update_status()

    datatracker = DataTracker(example)
    datatracker.start()

    while not quit:
        mainwin = example.get_win()
        #mainwin.refresh()

        c = mainwin.getkey()
        if c == 'q':
            datatracker.pause = True
            datatracker.quit = True
            while datatracker.isAlive():
                pass
            quit = True
        elif c == 'KEY_RESIZE' or  c == 'p':
            if c == 'KEY_RESIZE': example.resize()
            if example.status[1]<>1:
                datatracker.pause = True
                #mainwin.erase()
                #example.refresh_data(4.,4.)
                #example.update_fig()
                example.update_status(("*** Full view ***",1))
                datatracker.pause = False
            else:
                datatracker.pause = True
                example.update_status(("!!! Display has been paused. Press 'p' again to unpause !!!",-1))
        elif c == 'h':
            if example.status[1]<>2:
                datatracker.pause = True
                mainwin.erase()
                example.show_help()
                mainwin.refresh()
                example.update_status(("*** Help screen ***",2))
            else:
                datatracker.pause = False
                curses.ungetch('p')
        elif c == ' ':
            datatracker.pause = True
            example.update_status(("<<< Press a key to use as the style >>>",3))
            #mainwin.erase()
            global STYLE
            STYLE = mainwin.getkey()
            curses.ungetch('p')
            datatracker.pause = False
        elif c == 'd':
            datatracker.pause = True
            example.update_status(("<<< Enter string representing data parameter changes (a,b) >>>",4))
            curses.echo()
            para = [float(s) for s in (example.statuswin.getstr(1,2)).split(',')]
            curses.noecho()
            datatracker.para = para
            curses.ungetch('p')
            datatracker.pause = False
        else:
            example.update_status(("!!! %s key not understood !!!" %c,-1))

    #sleep(4)

if __name__ == "__main__":
    curses.wrapper(main)

